import gradio as gr

from audio.audio_processor import AudioProcessor
from debate_controller_backup import DebateController
import time
from audio.chattts import text_to_speech
from pydub import AudioSegment
import threading
import queue

# 在文件顶部添加CSS定义
css = """
.vertical-text{
    height: 550px !important;
    display: flex !important;
    flex-direction: column !important;
    gap: 20px !important;  /* 控制元素间距 */
}

/* 为包含HighlightedText的容器添加样式 */
.speakers-container {
    height: 100% !important;
    display: flex !important;
    flex-direction: column !important;
    justify-content: flex-end !important;
}
"""
js = """
                () => {
                    setInterval(() => {
                        const chatbot = document.querySelector('.gradio-chatbot');
                        if(chatbot) {
                            chatbot.scrollTop = chatbot.scrollHeight;
                        }
                    }, 100);
                }
                """


def main():
    with gr.Blocks(title="智能辩论系统", css=css, js=js) as demo:

        # region === 界面元素 ===
        # 在main函数开始处添加状态管理
        debate_state = gr.State({
            'running': False,
            'speaker': "主持人",
            'debate_state': [],  # 新增音频队列
            'audio_count': 0,  # 新增音频片段计数器
            'play_voice': False
        })
        # === 标题行 ===
        gr.Markdown("<h1 style='text-align: center;'>基于多智能体的模拟辩论系统</h1>")

        # === 配置行 ===
        with gr.Row():
            debate_mode = gr.Dropdown(
                label="辩论模式",
                choices=["智能体辩论", "人机辩论"],
                value="智能体辩论"
            )
            debate_topic = gr.Textbox(
                label="辩论题目",
                placeholder="请输入辩论主题...",
                value="女性更适合开出租车。",
            )
            agent_type = gr.Dropdown(
                label="辩论智能体",
                choices=["简单智能体", "反思智能体"],
                value="简单智能体"
            )

        # === 控制按钮 ===
        with gr.Row():
            start_btn = gr.Button("▶️ 开始")
            stop_btn = gr.Button("⏹️ 停止", interactive=False)
            reset_btn = gr.Button("🔄 重启")

        # === 主内容区 ===
        with gr.Row():
            # 左侧辩手面板
            with gr.Column(scale=1):
                pro_panel_label = gr.Markdown("### 正方")  # 添加变量引用
                pro_role = gr.Dropdown(
                    label="角色选择",
                    choices=["正方", "反方"],
                    value="正方"
                )
                with gr.Column(elem_classes="speakers-container"):
                    pro_speakers = gr.HighlightedText(
                        label="当前辩手",
                        elem_classes=["vertical-text"],
                        value=[("正方一辩", "蓄势待发"), ("正方二辩", "蓄势待发")],
                        color_map={"正在发言": "green", "蓄势待发": "gray"},
                        combine_adjacent=True,
                        adjacent_separator='\n\n'
                    )

                # 中央辩论大屏
            with gr.Column(scale=3):
                debate_screen = gr.Chatbot(
                    label="辩论实况",
                    height=630,
                    type='messages',
                    render_markdown=True,  # 启用Markdown渲染
                    bubble_full_width=False
                )
                is_play = gr.Checkbox(label="播放语音", value=False)
                soundfile = gr.Audio(
                    label="语音输出",
                    visible=False,
                    autoplay=True
                )
            # 右侧动态面板
            with gr.Column(scale=1, elem_classes="right-panel"):
                right_panel_label = gr.Markdown("### 反方")
                con_role = gr.Dropdown(
                    label="角色选择",
                    choices=["反方", "正方"],
                    value="反方"
                )
                with gr.Column():
                    con_speakers = gr.HighlightedText(
                        label="当前辩手",
                        elem_classes=["vertical-text"],
                        value=[("反方一辩", "蓄势待发"), ("反方二遍", "蓄势待发")],
                        color_map={"正在发言": "red", "蓄势待发": "gray"},
                        combine_adjacent=True,
                        adjacent_separator='\n\n'

                    )

                # 人机交互组件（默认隐藏）
                with gr.Column(visible=False) as human_interface:
                    human_input = gr.Textbox(
                        label="人类辩手输入",
                        lines=12
                    )
                    with gr.Row():
                        record_btn = gr.Audio(sources=["microphone"], streaming=True)  # 修改后的Audio组件
                        send_btn = gr.Button("发送")

        # endregion

        # region === 交互逻辑 ===
        def update_role(role):
            opposite = "反方" if role == "正方" else "正方"
            return [
                gr.update(value=opposite),  # 更新对方下拉框
                gr.update(value=f"### {role}"),  # 更新本方标题
                gr.update(value=f"### {opposite}")  # 更新对方标题
            ]

        def update_agent_list(agent):
            return gr.update(value=f"### {'人类' if agent == '人机辩论' else '反方'}")  # 修改这行

        def update_right_panel(mode):
            if mode == "人机辩论":
                return [
                    gr.update(visible=True),  # 显示语言输入等组件
                    gr.update(elem_classes=""),  # 缩小高亮小面板
                ]
            else:
                return [
                    gr.update(visible=False),  # 更新对方下拉框
                    gr.update(elem_classes=["vertical-text"]),
                ]

        # endregion

        # region 状态定义
        def reset_debate(state):
            state['running'] = False

            return [
                [],  # 对应debate_screen
                state,
                gr.update(interactive=True),  # start_btn
                gr.update(interactive=False),  # stop_btn
                gr.update(value=""),  # debate_topic
                gr.update(interactive=True),  # debate_mode
                gr.update(interactive=True),  # agent_type
                gr.update(interactive=True),  # pro_role
                gr.update(interactive=True),  # con_role
                gr.update(interactive=True),  # is_play
            ]

        def start_debate(state, play_voice):
            state['running'] = True
            state['play_voice'] = play_voice
            return [
                [],  # 对应debate_screen
                state,
                gr.update(interactive=False),  # start_btn
                gr.update(interactive=True),  # stop_btn
                gr.update(interactive=False),  # debate_topic
                gr.update(interactive=False),  # debate_mode
                gr.update(interactive=False),  # agent_type
                gr.update(interactive=False),  # pro_role
                gr.update(interactive=False),  # con_role
                gr.update(interactive=False)

            ]

        # endregion
        def update_interactive(state, speaker):
            for size in state['debate_state']:
                for debater_info in size:
                    if debater_info[0] == speaker:
                        debater_info[1] = "正在发言"
                    else:
                        debater_info[1] = "蓄势待发"

        def show_text(history, state, current_role, message):
            history.append({"role": current_role, "content": ""})
            for character in message:
                if state["running"] is False:
                    break
                history[-1]['content'] += character
                time.sleep(0.07)
                # yield history
                yield history, state, None, state['debate_state'][0], state['debate_state'][1]

        def show_audio(history, state, current_role, message):
            audio_processor = AudioProcessor(message)
            state['audio_count'] += audio_processor.audio_count  # 记录总音频段数
            # 立即更新聊天记录
            while state['audio_count'] > 0:
                audio_info, sleep_time = audio_processor.get_audio()
                if audio_info:
                    print(f"audio_info: {audio_info[0]}")
                    history.append({"role": current_role, "content": audio_info[2]})
                    yield history, state, audio_info[1], state['debate_state'][0], state['debate_state'][1]
                    state['audio_count'] -= 1
                    time.sleep(sleep_time)

        def debate_generator(topic, history: list, state: dict):
            topic = topic.strip()
            if topic == "":
                return [], state
            controller = DebateController(topic)
            state["debate_state"] = controller.get_debates_info()
            current_role = "assistant"  # 初始角色
            try:
                for speaker, message in controller.run_debate():
                    if state['running'] is False:
                        break
                    current_role = "user" if current_role == "assistant" else "assistant"
                    update_interactive(state, speaker)
                    playing = state['play_voice']
                    if playing:
                        yield from show_audio(history, state, current_role, message)
                    else:
                        yield from show_text(history, state, current_role, message)


            except Exception as e:
                yield history, state, None, None, None

        def stop_debate(state):
            state['running'] = False
            return [

                state,
                gr.update(interactive=True),  # start_btn
                gr.update(interactive=False),  # stop_btn
                gr.update(interactive=True),  # debate_topic
                gr.update(interactive=True),  # debate_mode
                gr.update(interactive=True),  # agent_type
                gr.update(interactive=True),  # pro_role
                gr.update(interactive=True)  # con_role
            ]

        # endregion

        # region 角色联动
        pro_role.change(
            update_role,
            pro_role,
            [con_role, pro_panel_label, right_panel_label]  # 添加两个标题输出
        )
        con_role.change(
            update_role,
            con_role,
            [pro_role, right_panel_label, pro_panel_label]  # 反向联动
        )

        # endregion

        # region 动态更新右侧面板
        debate_mode.change(
            update_right_panel,
            inputs=debate_mode,
            outputs=[human_interface, con_speakers]
        )

        agent_type.change(
            update_agent_list,
            agent_type,
            right_panel_label
        )

        is_play.change(
            lambda x: gr.update(visible=x),  # 使用 lambda 替代独立函数
            inputs=is_play,
            outputs=soundfile
        )

        # endregion

        # region 在控制按钮事件中更新状态
        start_btn.click(
            start_debate,
            inputs=[debate_state, is_play],
            outputs=[
                debate_screen,  # 清空聊天记录
                debate_state,  # 更新状态
                start_btn,  # 启用开始按钮
                stop_btn,  # 禁用停止按钮
                debate_topic,  # 清空辩论题目
                debate_mode,  # 启用辩论模式选择
                agent_type,  # 启用辩论智能体
                pro_role,  # 左侧角色选择
                con_role,
                is_play
                # 右侧角色选择
            ],
            queue=False
        ).then(
            debate_generator,
            inputs=[debate_topic, debate_screen, debate_state],
            outputs=[debate_screen, debate_state, soundfile, pro_speakers, con_speakers]
        )
        stop_btn.click(
            stop_debate,
            inputs=debate_state,
            outputs=[
                debate_state,
                start_btn,
                stop_btn,
                debate_topic,
                debate_mode,
                agent_type,
                pro_role,  # 左侧角色选择
                con_role]
        )
        reset_btn.click(
            reset_debate,
            inputs=debate_state,
            outputs=[
                debate_screen,  # 清空聊天记录
                debate_state,  # 更新状态
                start_btn,  # 启用开始按钮
                stop_btn,  # 禁用停止按钮
                debate_topic,  # 清空辩论题目
                debate_mode,  # 启用辩论模式选择
                agent_type,  # 启用辩论智能体
                pro_role,  # 左侧角色选择
                con_role,  # 右侧角色选择
                is_play
            ]
        )

    # endregion

    demo.launch()


if __name__ == "__main__":
    main()
